import itertools
import numpy as np
import pandas as pd

# HMM parameters
A = np.array([
    [0.5, 0.2, 0.3],
    [0.2, 0.4, 0.4],
    [0.4, 0.1, 0.5]
])

B = np.array([
    [0.8, 0.2],  # state 0: P(0), P(1)
    [0.1, 0.9],  # state 1
    [0.5, 0.5]   # state 2
])

pi = np.array([0.5, 0.3, 0.2])  # Initial state distribution

# Observation sequence: 0 1 0 1
obs = [0, 1, 0, 1]
n_states = 3
sequence_length = len(obs)

# Enumerate all possible state sequences (3^4 = 81)
state_sequences = list(itertools.product(range(n_states), repeat=sequence_length))

# Calculate probabilities
results = []
total_posterior = 0.0
for seq in state_sequences:
    # Prior probability
    prior = pi[seq[0]]
    for t in range(1, sequence_length):
        prior *= A[seq[t - 1], seq[t]]

    # Likelihood
    likelihood = 1.0
    for t, o in enumerate(obs):
        likelihood *= B[seq[t], o]

    # Posterior (unnormalized)
    posterior = prior * likelihood
    results.append({
        'seq': seq,
        'prior': prior,
        'likelihood': likelihood,
        'posterior': posterior
    })
    total_posterior += posterior

# Normalize posterior
for r in results:
    r['posterior'] /= total_posterior

# Sort by posterior probability
top3 = sorted(results, key=lambda x: -x['posterior'])[:3]

# Prepare output table
table_data = []
for r in top3:
    table_data.append([
        ''.join(map(str, r['seq'])),
        round(r['prior'], 6),
        round(r['likelihood'], 6),
        round(r['posterior'], 6)
    ])

df = pd.DataFrame(table_data, columns=[
    "Most Probable Hidden State Sequences",
    "Prior Probability",
    "Likelihood",
    "Posterior Probability"
])

print(df)